import os
import json
import cv2
import numpy as np
import torch
from PIL import Image, ImageDraw, ImageFont
from torchvision import transforms
import matplotlib.pyplot as plt

from model import densenet121
plt.rcParams['font.sans-serif'] = ['SimHei']


def cv2AddChineseText(img, text, position, textColor, textSize):
    if (isinstance(img, np.ndarray)):  # 判断是否OpenCV图片类型
        img = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
    # 创建一个可以在给定图像上绘图的对象
    draw = ImageDraw.Draw(img)
    # 字体的格式
    fontStyle = ImageFont.truetype(
        "./simsun.ttc", textSize, encoding="utf-8")
    # 绘制文本
    draw.text(position, text, textColor, font=fontStyle)
    # 转换回OpenCV格式
    return cv2.cvtColor(np.asarray(img), cv2.COLOR_RGB2BGR)


def yuce_huwai(root):
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

    data_transform = transforms.Compose(
        [transforms.Resize(256),
         transforms.CenterCrop(224),
         transforms.ToTensor(),
         transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])])

    # load image
    img_path = root
    assert os.path.exists(img_path), "file: '{}' dose not exist.".format(img_path)
    img_or = Image.open(img_path)
    # [N, C, H, W]
    img = data_transform(img_or)
    # expand batch dimension
    img = torch.unsqueeze(img, dim=0)

    # read class_indict
    json_path = './class_indices.json'
    assert os.path.exists(json_path), "file: '{}' dose not exist.".format(json_path)

    with open(json_path, "r") as f:
        class_indict = json.load(f)

    # create model
    model = densenet121(num_classes=2).to(device)
    # load model weights
    model_weight_path = "./weights/model-39.pth"
    model.load_state_dict(torch.load(model_weight_path, map_location=device))
    model.eval()
    with torch.no_grad():
        # predict class
        output = torch.squeeze(model(img.to(device))).cpu()
        predict = torch.softmax(output, dim=0)
        predict_cla = torch.argmax(predict).numpy()

    if class_indict[str(predict_cla)]=='meiquan':
        text = '未苫盖'
        img_ori = cv2.imread(img_path)
        img_A = cv2AddChineseText(img_ori, text, (200, 200), (254, 0, 0), 200)
        cv2.imwrite('./cap_huwai/1.jpg', img_A)
    else:
        text = '苫盖'
        img_ori = cv2.imread(img_path)
        img_A = cv2AddChineseText(img_ori, text, (200, 200), (254, 0, 0), 200)
        cv2.imwrite('./cap_huwai/1.jpg', img_A)

    a = "class: {:10}".format(class_indict[str(0)])
    a_prob = "prob: {:.3}".format(predict[0].numpy())
    b = "class: {:10}".format(class_indict[str(1)])
    b_prob = "prob: {:.3}".format(predict[1].numpy())
    liebiao_a = [a,a_prob]
    liebiao_b = [b,b_prob]
    print(liebiao_a,liebiao_b)
    return liebiao_a,liebiao_b